package wangn;

/*
* 作者：王宁 软件学院 161250139
* 时间：2018/10/25 完成
*      2018/12/6 修复bug且加入教师检测程序
*
* 红黑树 Red-Black Tree simple implementation
*
* 用法：
* ```
* wangn.RedBlackTree redBlackTree = new wangn.RedBlackTree();
* redBlackTree.insert(integer);
* redBlackTree.delete(integer);
* boolean isKeyExist = redBlackTree.find(integer);
* String sequenceFormatStr = redBlackTree.toString();
* ```
*
* 注意：
* 1、树的节点没有使用父指针，所以一旦有向上调整时候都会搜索父节点位置，故本算法实现效率不算太高
*
* */
public class RedBlackTree {

    private final static int RED = 0;
    private final static int BLACK = 1;

    /*
     * class RBTNode    RBT节点
     *    key : 键
     *    color ：颜色
     *    left ： 左孩子
     *    right ： 右孩子
     */
    class RBTNode{
        int key;
        int color;
        RBTNode left;
        RBTNode right;
        public RBTNode(int k){
            this.key = k;
            this.color = RED;                                           //初始化 color = RED
        }
    }

    //root of red-black tree
    private RBTNode root;
    private static boolean isRootValid = false;                          //当前root是否存在

    /*提供给外面调用的接口
     * @param key 想要插入的值
     * @return 新的root
     * */
    public RBTNode insert(int key){
        if (!isRootValid){
            root = new RBTNode(key);
            isRootValid = true;
        }else{
            InsReturn insReturn = insert_rec(root, key);
            root = insReturn.newTree;                                   //root可能被更改
        }
        root.color = BLACK;                                             //ARB树改成RB树
        return root;
    }

    /*提供给外面调用的接口
     * @param key 想要删除的值
     * @return 新的root或者null
     * */
    public void delete(int key){
        if (!isRootValid)return;                                          //树不存在就返回null

        isNeedUpFix=false;

        RBTNode fatherNode;
        if (root.key==key)fatherNode=null;                              //此时不存在父节点，但是后面调整的时候可能会更改
        else {
            fatherNode = findFather(root,key);
            if (fatherNode==null)return;                                //此时找不到这个节点
        }

        //找到待删除的节点
        RBTNode toBeDelNode = fatherNode==null?root:(fatherNode.left==null?(fatherNode.right):(fatherNode.left.key==key?(fatherNode.left):(fatherNode.right)));

        if (toBeDelNode.right != null){                                 //若可以找到后继节点
            RBTNode tmp = toBeDelNode.right;                            //找到这个后继节点并且用后继节点的值替换当前要删除的节点，然后要删除的节点转移到后继节点上
            RBTNode father = toBeDelNode;                               //后继节点的父节点
            RBTNode grandFather = fatherNode;
            while (tmp.left != null){
                grandFather = father;
                father = tmp;
                tmp = tmp.left;
            }
            toBeDelNode.key = tmp.key;
            if (father == toBeDelNode){
                father.right = tmp.right;                               //删除tmp
            }else{
                father.left = tmp.right;
            }

            if (tmp.color == RED){                                      //若删除的是红节点就直接返回
                return;
            }
            if (tmp.right != null){                                     //若补上来的tmp的右子树非null，即右子树是红色的，那么将其变为黑色，就可以了。
                tmp.right.color = BLACK;
                return;
            }

            RBTNode newFatherNode;
            if (father == toBeDelNode){                                 //若待删除节点的后继节点（真正删除的节点tmp）就是他的右子树，则其后继节点的右子树（可能是null或者红节点）作为右子树挂在待删除节点上
                 newFatherNode = fixRightAfterDelete(father, null);
            }
            else                                                        //否则后继节点作为左子树挂在它的父节点上
                newFatherNode = fixLeftAfterDelete(father, null);
            fixUp(grandFather, newFatherNode,father);
        }
        else if (toBeDelNode.left != null){                             //否则尝试找前驱节点，下面都和上面一样，但是结构相反
            RBTNode tmp = toBeDelNode.left;                             //找到这个前驱节点并且用前驱节点的值替换当前要删除的节点，然后要删除的节点转移到前驱节点上
            RBTNode father = toBeDelNode;                               //前驱节点的父节点
            RBTNode grandFather = fatherNode;
            while (tmp.right != null){
                grandFather = father;
                father = tmp;
                tmp = tmp.right;
            }
            toBeDelNode.key = tmp.key;
            if (father == toBeDelNode){
                father.left = tmp.left;                                 //删除tmp
            }else{
                father.right = tmp.left;
            }

            if (tmp.color == RED){                                      //若删除的是红节点就直接返回
                return;
            }
            if (tmp.left != null){                                      //若补上来的tmp的右子树非null，即左子树是红色的，那么将其变为黑色，就可以了。
                tmp.left.color = BLACK;
                return;
            }
            RBTNode newFatherNode;
            if (father == toBeDelNode){                                 //若待删除节点的前驱节点（真正删除的节点tmp）就是他的左子树，则其前驱节点的左子树（可能是null或者红节点）作为左子树挂在待删除节点上
                newFatherNode = fixLeftAfterDelete(father, null);
            }
            else newFatherNode = fixRightAfterDelete(father, null);          //否则前驱节点作为右子树挂在它的父节点上
            fixUp(grandFather, newFatherNode,father);

        }
        else{                                                           //此时待删除节点不能被替代，需要根据父节点将之删除，然后从父节点开始调整（上溯）
            if (fatherNode == null){                                    //若要删除的是root，而且root所在树只有root一个内节点，使得root无效即可
                isRootValid = false;
                return;
            }
            RBTNode newFatherNode;
            if (fatherNode.left == toBeDelNode){                        //作为左子树
                fatherNode.left = null;
                if (toBeDelNode.color == BLACK){                        //如果是黑色需要进行调整
                    newFatherNode = fixLeftAfterDelete(fatherNode, null);
                }else return;;
            }else{                                                      //作为右子树
                fatherNode.right = null;
                if (toBeDelNode.color == BLACK){
                    newFatherNode = fixRightAfterDelete(fatherNode, null);
                }else return;
            }
            RBTNode grandFather = findFather(root,fatherNode.key);
            fixUp(grandFather, newFatherNode, fatherNode);
        }
    }

    /*
    * @function 挂载新树和可能的向上修复
    * @param grandFather 子树的祖父节点
    * @param newFatherNode 新的子树父节点（子树的根节点）
    * @param father 以前的父节点 用于了解原来的子树所在位置
    * */
    private void fixUp(RBTNode grandFather, RBTNode newFatherNode, RBTNode father){
        if (grandFather == null)root = newFatherNode;               //开始修复，由于可能需要上溯，所以需要嵌套一次
        else{
            RBTNode anotherNode;                                    //fix后新的父节点
            RBTNode ancientNode;                                    //需要上溯的时候，需要先计算出grandFather的父节点
            if (grandFather.left == father){
                grandFather.left = newFatherNode;
            }
            else{
                grandFather.right = newFatherNode;
            }
            while (isNeedUpFix){
                isNeedUpFix = false;
                ancientNode = findFather(root, grandFather.key);
                if (grandFather.left == newFatherNode)
                    anotherNode = fixLeftAfterDelete(grandFather,newFatherNode);
                else
                    anotherNode = fixRightAfterDelete(grandFather,newFatherNode);
                if (ancientNode == null){
                    root = anotherNode;
                    isNeedUpFix = false;
                    return;
                }
                else {
                    if (ancientNode.left == grandFather)
                        ancientNode.left = anotherNode;
                    else
                        ancientNode.right = anotherNode;
                    newFatherNode = grandFather;
                    grandFather = ancientNode;
                }
            }
        }
    }

    /*
    * @function 寻找一个节点父亲  要求传入节点不能为null
    * @param oldTree 当前子树的根节点
    * @param key 待查找的值
    * @return RBTNode 待查找节点的父节点
    * */
    private RBTNode findFather(RBTNode oldTree, int key){
        if (oldTree.left != null ){
            if (oldTree.left.key == key)return oldTree;
            RBTNode leftNode = findFather(oldTree.left,key);
            if (leftNode != null)return leftNode;
        }
        if (oldTree.right != null){
            if (oldTree.right.key == key)return oldTree;
            RBTNode rightNode = findFather(oldTree.right,key);
            if (rightNode != null)return rightNode;
        }
        return null;
    }

    /*提供给外面调用的接口
     * @param key 想要查找的值
     * @return key是否在红黑树中
     * */
    public boolean find(int key){
        RBTNode findNode = find_rec(root, key);
        return findNode!=null;
    }

    /*
    * @function 查找元素，方法与普通二叉搜索树的search一样
    * @return 查到的节点或者null
    * */
    private RBTNode find_rec(RBTNode oldTree, int key){
        if (oldTree == null)return null;
        if (key==oldTree.key)return oldTree;
        if (key<oldTree.key){
            return find_rec(oldTree.left,key);
        } else{
            return find_rec(oldTree.right,key);
        }
    }

    private static StringBuilder toStringBuffer;

    private void traverseAndPrint(RBTNode tree){
        if (tree==null)return;
        traverseAndPrint(tree.left);
        toStringBuffer.append(tree.key+",");
        traverseAndPrint(tree.right);
    }

    //********************************内部实现**********************************//
    //*****************************红黑树的插入部分*******************************//

    private static final int Ok = 2;
    private static final int RBR = 3;
    private static final int BRB = 4;
    private static final int BRR = 5;
    private static final int RRB = 6;

    /*
    * InsReturn
    *   newTree : 新插入节点的子树
    *   status ： 子树的状态
    * */
    class InsReturn{
        RBTNode newTree;
        int status;
        public InsReturn(RBTNode tree){
            this.newTree = tree;
            this.status = BRB;                                          //初始状态是 BRB 形式的，为了清楚表示，后面都需要继续明确赋值
        }
    }

    /*
     * @function 插入新节点并且自下向上调整
    * @param:  oldTree 父节点
    * @return:  insReturn 包含子树的根和状态的对象
    * */
    private InsReturn insert_rec(RBTNode oldTree, int key){
        InsReturn insReturn = new InsReturn(oldTree);
        if (oldTree==null){
            insReturn.newTree = new RBTNode(key);
            insReturn.status = BRB;
        }
        else if (key<oldTree.key){
            InsReturn leftInsReturn = insert_rec(oldTree.left,key);
            insReturn = repairLeft(oldTree, leftInsReturn);             //修复左子树
        }
        else {
            InsReturn rightInsReturn = insert_rec(oldTree.right,key);
            insReturn = repairRight(oldTree, rightInsReturn);           //修复右子树
        }
        return insReturn;
    }

    /*
    * @function : 修复左子树
    * @param oldTree 父节点
    * @param leftInsReturn oldTree的左子树的根和状态信息
    * @return ： 新的子树的根和状态信息
    * */
    private InsReturn repairLeft(RBTNode oldTree, InsReturn leftInsReturn){
        InsReturn insReturn = new InsReturn(oldTree);
        oldTree.left = leftInsReturn.newTree;                           //链接左子树
        switch (leftInsReturn.status){
            case Ok:                                                    //不需要做改变
                insReturn.status = Ok;
                break;
            case RBR:                                                   //不需要额外调整
                insReturn.status = Ok;
                break;
            case BRB:                                                   //如果父亲是红色的，记录为RRB后向上调整，否则是黑色的不需要调整
                if (oldTree.color == RED){
                    insReturn.status = RRB;                             //因为在左子树，则为RRB
                }else {
                    insReturn.status = RBR;
                }
                break;
            default:                                                    //若兄弟节点是黑色的，则情况属于RRB或者BRR，使用rebalLeft处理；若兄弟节点是红色的，使用colorFlip改变颜色即可
                if (oldTree.right==null || oldTree.right.color==BLACK){
                    insReturn = rebalLeft(oldTree, leftInsReturn.status);
                    insReturn.status = RBR;                             //修复以后的状态是RBR
                }else{
                    colorFlip(oldTree);                                 //修复后状态一定是RBR
                    insReturn.status = BRB;
                }
        }
        return insReturn;
    }

    /*
    * @function 处理左子树中RRB或者BRR两种情况的
    * @param oldTree 左子树
    * @param status 左子树是BRR还是RRB
    * @return 新的子树的根和状态信息
    * */
    private InsReturn rebalLeft(RBTNode oldTree, int status){
        InsReturn insReturn = new InsReturn(oldTree);
        if (status==BRR){                                               //BRR型，B/R表示颜色，括号中是节点编号
            RBTNode L = oldTree;                                        //      B(1)                    B(3)
            RBTNode S = L.left;                                         //  R(2)           --->    R(2)       R(1)
            RBTNode M = S.right;                                        //      R(3)
            S.right = M.left;
            L.left = M.right;
            M.left = S;
            M.right = L;
            L.color = RED;
            M.color = BLACK;
            insReturn.newTree = M;
        }else{                                                          //RRB型，B/R表示颜色，括号中是节点编号
            RBTNode L = oldTree;                                        //          B(1)
            RBTNode M = L.left;                                         //      R(2)        --->          B(2)
            RBTNode S = M.left;                                         // R(3)                     R(3)        R(1)
            L.left = M.right;
            M.right = L;
            L.color = RED;
            M.color = BLACK;
            insReturn.newTree = M;
        }
        return insReturn;
    }

    /*
    * @function : 4-node的情况，通过更改颜色即可让原来的红黑树符合约束
    * */
    private void colorFlip(RBTNode oldTree){
        oldTree.color = RED;                                            //       B(1)                         R(1)
        oldTree.left.color = BLACK;                                     // R(2)         R(3)     --->   B(2)        B(3)
        oldTree.right.color = BLACK;
    }

    /*
    * @function 与 @code{repairLeft} 效果一样，只不过修复的是对称的情况
    * @param oldTree 父节点
    * @param rightInsReturn oldTree的右子树的根和状态信息
    * @return ： 新的子树的根和状态信息
    * */
    private InsReturn repairRight(RBTNode oldTree, InsReturn rightInsReturn){
        InsReturn insReturn = new InsReturn(oldTree);
        oldTree.right = rightInsReturn.newTree;                         //链接右子树
        switch (rightInsReturn.status){
            case Ok:                                                    //不需要做改变
                insReturn.status = Ok;
                break;
            case RBR:                                                   //不需要额外调整
                insReturn.status = Ok;
                break;
            case BRB:                                                   //如果父亲是红色的，记录为RRB后向上调整，否则是黑色的不需要调整
                if (oldTree.color == RED){
                    insReturn.status = BRR;                             //因为在左子树，则为BRR
                }else {
                    insReturn.status = RBR;
                }
                break;
            default:                                                    //若兄弟节点是黑色的，则情况属于RRB或者BRR，使用rebalRight处理；若兄弟节点是红色的，使用colorFlip改变颜色即可
                if (oldTree.left==null || oldTree.left.color==BLACK){
                    insReturn = rebalRight(oldTree, rightInsReturn.status);
                    insReturn.status = RBR;                             //修复以后的状态是RBR
                }else{
                    colorFlip(oldTree);                                 //修复后状态一定是RBR
                    insReturn.status = BRB;
                }
        }
        return insReturn;
    }

    /*
    * @function 与@code{rebalLeft}功能一样，修复的是右子树
    * @param oldTree 右子树
    * @param status 右子树是BRR还是RRB
    * @return 新的子树的根和状态信息
    * */
    private InsReturn rebalRight(RBTNode oldTree, int status){
        InsReturn insReturn = new InsReturn(oldTree);
        if (status == RRB){                                             //RRB型，B/R表示颜色，括号中是节点编号
            RBTNode S = oldTree;                                        //B(1)                        B(3)
            RBTNode L = S.right;                                        //        R(2)    --->   R(1)       R(2)
            RBTNode M = L.left;                                         //    R(3)
            S.right = M.left;
            L.left = M.right;
            M.left = S;
            M.right = L;
            M.color = BLACK;
            S.color = RED;
            insReturn.newTree = M;
        }else{                                                          //BRR型号，B/R表示颜色，括号中是节点编号
            RBTNode S = oldTree;                                        //B(1)
            RBTNode M = S.right;                                        //      R(2)        ---->       B(2)
            RBTNode L = M.right;                                        //          R(3)            R(1)    R(3)
            S.right = M.left;
            M.left = S;
            M.right = L;
            M.color = BLACK;
            S.color = RED;
            insReturn.newTree = M;
        }
        return insReturn;
    }


    //*********************************内部实现*******************************//
    //*****************************红黑树的删除部分*******************************//

    private static boolean isNeedUpFix = false;                                                 //约定，如果出现需要上溯修复的时候，下面的方法设置为true，由调用者处理

    /*
     * @function 修复作为左子树的节点被删除的情况，此时删除的是黑节点，替代的（targetNode）是黑节点或者null（看成黑色）
    * @param fatherNode 父节点
    * @param targetNode 变换规则的目标节点
    * @return RBTNode 返回fatherNode这棵树调整后的新的父节点
    * */
    private RBTNode fixLeftAfterDelete(RBTNode fatherNode, RBTNode targetNode){
        RBTNode P = fatherNode;
        RBTNode D = targetNode;
        RBTNode S = P.right;                                                                    //S与原来删除的节点是同高的，因此S不可能为null
        RBTNode SL = S.left;
        RBTNode SR = S.right;

        //case1 : 兄弟节点是红色的   此时保证SL、SR也不可能是null
        if (S.color == RED){                                                                    //颜色如代码描述
            P.right = SL;                                                                       //      P                           S
            S.left = P;                                                                         //D           S         --->    P       SR
            P.color = RED;                                                                      //          SL  SR          D     SL
            S.color = BLACK;
            S.left = fixLeftAfterDelete(P,D);                                                   //此时D与SL高度不平衡，需要继续向下调整
            return S;
        }
        //case2 : SL红色
        else if (SL != null && SL.color == RED){
            P.right = SL.left;                                                                  //        P                         SL
            S.left = SL.right;                                                                  //   D          S       --->    P      S
            SL.right = S;                                                                       //            SL  SR          D          SR
            SL.left = P;
            SL.color = P.color;
            P.color = BLACK;
            return SL;
        }
        //case3: SR红色
        else if (SR != null && SR.color == RED){
            P.right = SL;                                                                       //      P                           S
            S.left = P;                                                                         //D             S       --->    P       SR
            S.color = P.color;                                                                  //          SL     SR         D    SL
            P.color = BLACK;
            SR.color = BLACK;
            return S;
        }
        //case4 : P红色
        else if (P.color == RED){
            P.color = BLACK;
            S.color = RED;
            return P;
        }
        //case5 : P黑色，需要向上调整
        else if (P.color == BLACK){
            S.color = RED;
            isNeedUpFix = true;
            return P;
        }
        return fatherNode;
    }

    /*
     * @function 修复作为右子树的节点被删除的情况，此时删除的是黑节点，替代的是黑节点或者null（看成黑色）
     * @param fatherNode 父节点
     * @param targetNode 变换规则的目标节点
     * @return RBTNode 返回fatherNode这棵树调整后的新的父节点
     * */
    private RBTNode fixRightAfterDelete(RBTNode fatherNode, RBTNode targetNode){
        RBTNode P = fatherNode;
        RBTNode D = targetNode;
        RBTNode S = P.left;                                                                    //S与原来删除的节点是同高的，因此S不可能为null
        RBTNode SL = S.left;
        RBTNode SR = S.right;

        //case1 : 兄弟节点是红色的   此时保证SL、SR也不可能是null
        if (S.color == RED){                                                                    //颜色如代码描述
            P.left = SR;                                                                        //      P                    S
            S.right = P;                                                                        //   S     D     --->    SL       P
            P.color = RED;                                                                      // SL  SR                      SR    D
            S.color = BLACK;
            S.right = fixRightAfterDelete(P,D);                                                   //此时D与SL高度不平衡，需要继续向下调整
            return S;
        }
        //case2 : SR红色
        else if (SR != null && SR.color == RED){
            P.left = SR.right;                                                                  //        P                     SR
            S.right = SR.left;                                                                  //     S     D    --->        S    P
            SR.left = S;                                                                        //   SL  SR                SL         D
            SR.right = P;             //
            SR.color = P.color;
            P.color = BLACK;
            return SR;
        }
        //case3: SL红色
        else if (SL != null && SL.color == RED){
            P.left = SR;                                                                         //        P                           S
            S.right = P;                                                                         //     S      D   --->           SL       P
            S.color = P.color;                                                                   //   SL  SR                             SR  D
            P.color = BLACK;
            SL.color = BLACK;
            return S;
        }
        //case4 : P红色
        else if (P.color == RED){
            P.color = BLACK;
            S.color = RED;
            return P;
        }
        //case5 : P黑色，需要向上调整
        else if (P.color == BLACK){
            S.color = RED;
            isNeedUpFix = true;
            return P;
        }
        return fatherNode;
    }

    /***********************************************************/
    /*下面是把老师的检测程序放过来做了一点改动（非逻辑部分）*/
    /***********************************************************/

    /**
     * @return 将红黑树以测试要求的格式输出
     */
    @Override
    public String toString(){
        return getString(this.root);
    }

    private String getString(RBTNode rbtNode)
    {
        String s = "(";
        if(rbtNode.left == null)
            s += "NIL";
        else
            s += getString(rbtNode.left);
        s += ", (";
        s += "" + rbtNode.key + "," + (rbtNode.color == RED ? "red" : "black") + "), ";

        if(rbtNode.right == null)
            s += "NIL";
        else
            s += getString(rbtNode.right);

        s += ")";

        return s;
    }

    /**
     * @return 是否满足颜色约束
     */
    public boolean checkColor(){
        return checkColor_rec(this.root);
    }

    private boolean checkColor_rec(RBTNode rbtNode)
    {
        int leftColor;
        if(rbtNode.left == null)
            leftColor = BLACK;
        else
        {
            if(!checkColor_rec(rbtNode.left))
                return false;
            leftColor = rbtNode.left.color;
        }
        int rightColor;
        if(rbtNode.right == null)
            rightColor = BLACK;
        else
        {
            if(!checkColor_rec(rbtNode.right))
                return false;
            rightColor = rbtNode.right.color;
        }

        if(rbtNode.color == BLACK)
            return true;
        if(leftColor == RED || rightColor == RED)
            return false;
        return true;
    }

    /**
     * @return n如果树的满足各外部节点的黑路径长度相等，否则-1
     */
    public int getBlackHeight(){
        return getBlackHeight_rec(this.root);
    }

    private int getBlackHeight_rec(RBTNode rbtNode)
    {
        int lHeight;
        if(rbtNode.left == null)
            lHeight = 1;
        else
        {
            lHeight = getBlackHeight_rec(rbtNode.left);
            if(lHeight < 0)
                return -1;
        }

        int rHeight;
        if(rbtNode.right == null)
            rHeight = 1;
        else
        {
            rHeight = getBlackHeight_rec(rbtNode.right);
            if(rHeight < 0)
                return -1;
        }

        if(lHeight == rHeight)
            return lHeight + (rbtNode.color == BLACK ? 1 : 0);

        return -1;
    }

}
